import time
import streamlit as st
import pandas as pd
import logging
import os
from datetime import timedelta
from model_calls import (
    generate_sql_cached,
    should_generate_chart_cached,
    generate_plotly_code_cached,
    generate_plot_cached,
    run_query_cached
)

avatar_url = "https://vanna.ai/img/vanna.svg"

# Ensure log directory exists
log_dir = "logs"
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

# Initialize logger
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# Add handlers if not already present
if not logger.handlers:
    # File handler
    file_handler = logging.FileHandler(os.path.join(log_dir, 'app.log'))
    file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
    logger.addHandler(file_handler)
    
    # Console handler
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
    logger.addHandler(console_handler)

# Force flush logs
def flush_logs():
    for handler in logger.handlers:
        handler.flush()

# Configure Streamlit
st.set_page_config(
    page_title="Data Analytics",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Increase timeout
if not st.session_state.get("configured"):
    st.session_state["server.maxMessageSize"] = 1000
    st.session_state["server.timeout"] = 300
    st.session_state["configured"] = True

# Force log flush after each operation
flush_logs()

st.sidebar.title("Output Settings")
st.sidebar.checkbox("Show Params", value=True, key="show_params")
st.sidebar.checkbox("Show Table", value=True, key="show_table")
st.sidebar.checkbox("Show Plotly Code", value=True, key="show_plotly_code")
st.sidebar.checkbox("Show Chart", value=True, key="show_chart")
st.sidebar.checkbox("Show Summary", value=True, key="show_summary")
st.sidebar.checkbox("Show Follow-up Questions", value=True, key="show_followup")
st.sidebar.button("Reset", on_click=lambda: set_question(None), use_container_width=True)

st.title("Vanna AI")
# st.sidebar.write(st.session_state)

def set_question(question):
    st.session_state["my_question"] = question

assistant_message_suggested = st.chat_message(
    "assistant", avatar=avatar_url
)
if assistant_message_suggested.button("Click to show suggested questions"):
    st.session_state["my_question"] = None
    questions = ['What is the average age of people who live in New York?', 'How many people live in New York?', 'What is the average income of people who live in New York?']
    for i, question in enumerate(questions):
        time.sleep(0.05)
        button = st.button(
            question,
            on_click=set_question,
            args=(question,),
        )

my_question = st.session_state.get("my_question", default=None)

if my_question is None:
    my_question = st.chat_input(
        "Ask me a question about your data",
    )

if my_question:
    st.session_state["my_question"] = my_question
    user_message = st.chat_message("user")
    user_message.write(f"{my_question}")

    params = generate_sql_cached(question=my_question)

    if params:
        if st.session_state.get("show_params", True):
            assistant_message_sql = st.chat_message(
                "assistant", avatar=avatar_url
            )
            assistant_message_sql.code(params, language="json", line_numbers=True)

        # df = pd.DataFrame({"city": ["New York", "New York", "New York"], "age": [25, 30, 35], "income": [50000, 60000, 70000]})
        df = run_query_cached(params=params)

        if df is not None:
            st.session_state["df"] = df

        if st.session_state.get("df") is not None:
            if st.session_state.get("show_table", True):
                df = st.session_state.get("df")
                assistant_message_table = st.chat_message(
                    "assistant",
                    avatar=avatar_url,
                )
                if len(df) > 10:
                    assistant_message_table.text("First 10 rows of data")
                    assistant_message_table.dataframe(df.head(10))
                else:
                    assistant_message_table.dataframe(df)

            if should_generate_chart_cached(question=my_question, params=params, df=df):

                code = generate_plotly_code_cached(question=my_question, params=params, df=df)

                if st.session_state.get("show_plotly_code", False):
                    assistant_message_plotly_code = st.chat_message(
                        "assistant",
                        avatar=avatar_url,
                    )
                    assistant_message_plotly_code.code(
                        code, language="python", line_numbers=True
                    )

                if code is not None and code != "":
                    if st.session_state.get("show_chart", True):
                        assistant_message_chart = st.chat_message(
                            "assistant",
                            avatar=avatar_url,
                        )
                        fig = generate_plot_cached(code=code, df=df)
                        if fig is not None:
                            assistant_message_chart.plotly_chart(fig)
                        else:
                            assistant_message_chart.error("I couldn't generate a chart")
    else:
        assistant_message_error = st.chat_message(
            "assistant", avatar=avatar_url
        )
        assistant_message_error.error("I wasn't able to generate Params for that question")